import typing

import numpy
import numpy as np
import scipy.constants as C
import scipy.interpolate
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt
from scipy.special import j1, jn_zeros
from scipy.interpolate import interp1d
from scipy.integrate import trapezoid
import common
import matplotlib
from  _logging import  logger
matplotlib.use('tkagg')
# 电子运动方程参数
f0 = 9.3e9  # 电场频率 (Hz)
v_0 = 2.85e7  # 平均速度 (m/s)
delta_v = 1e7  # 速度振幅 (m/s)
T_0 = 1 / f0  # 射频周期 (s)
omega_0 = 2 * np.pi * f0  # 角频率

# 空间电荷场参数
Q = 1e-9   # 单个电子圆盘带电量
a = 5e-3
b = 4e-3
N_zeros = 70

# 腔场参数
R_Q = 128  # 腔的R/Q
Q_t = 298  # 品质因子
omega_cav = 2 * np.pi * 3.45e9  # 腔的谐振频率 (Hz)
v_e = 0.5 * C.c  # 电子速度 (m/s)
M = 0.688  # 耦合系数

from field_generator import space_charge
sc_generator = space_charge.RingSpaceChaegeFieldGenerator(Q, a,b)


# 电场空间形状函数 f(z)
def f(z, k=0.5 * (2*numpy.pi / (v_e / f0)) / np.sqrt(-np.log(M))):
    return (k / np.sqrt(np.pi)) * np.exp(-k ** 2 * (z) ** 2)

zs = numpy. array([-0.02330325, -0.02235173, -0.0214002 , -0.02044868, -0.01949716,
       -0.01854564, -0.01759412, -0.0166426 , -0.01569108, -0.01473956,
       -0.01378803, -0.01283651, -0.01188499, -0.01093347, -0.00998195,
       -0.00903043, -0.00807891, -0.00712738, -0.00617586, -0.00522434,
       -0.00427282, -0.0033213 , -0.00236978, -0.00141826, -0.00046673,
        0.00048479,  0.00143631,  0.00238783,  0.00333935,  0.00429087,
        0.00524239,  0.00619391,  0.00714544,  0.00809696,  0.00904848,
        0.01      ])
f(zs)
# 电子圆盘的速度 (时间变化)
def v_i(t, initial_speed=v_0, delta_v=delta_v):
    return initial_speed + delta_v * np.cos(omega_0 * t)


# 电子运动方程中的洛伦兹力计算
def lorentz(t, y, z_i, Q, a, b,Vcav,
            # Vcav_history:typing.List[float],
            zmin
            ):
    z, pz = y[0::2], y[1::2]
    gamma_factor = np.sqrt(pz ** 2 + (C.m_e * C.c) ** 2)
    dz_dt = pz * C.c / gamma_factor  # dz/dt
    # E_spch_vals = E_spch(z, z_i, Q, a, b)
    # E_cav_vals = E_cav(t, z, z_i,Vcav, )
    # Ez = E_spch_vals + E_cav_vals
    Ez = get_E(t, z, z_i,Vcav,zmin)

    dpz_dt = -Q * Ez  # dpz/dt (Lorentz force)
    return np.stack([dz_dt, dpz_dt], axis=1).flatten()


# 电场函数 (空间电荷场)
def Es(z, Q, a, b):
    mu_0_ps =np.matrix( jn_zeros(0, N_zeros).reshape((-1, 1)))
    arr = np.exp(-mu_0_ps * np.abs(z) / a).A * (np.matrix((2 / mu_0_ps.A * j1(mu_0_ps.A * b / a) / j1(mu_0_ps.A)) ** 2 )* np.sign(z)).A
    return Q / (2 * np.pi * C.epsilon_0 * b ** 2) * arr.sum(axis = 0)


# 总的空间电荷场 E_spch (z,Espch_z)
def E_spch(z, z_i, Q, a, b):
    E_spch_total = np.zeros_like(z)
    for z_pos in z_i:
        E_spch_total += Es(z - z_pos, Q, a, b)
    return E_spch_total


# 腔场 E_cav 计算
def E_cav(t,    z,   Vtg :complex=0j):
    # velocities = v_i(t)
    # I_ind = np.sum(Q_i * velocities * f(z_i))  # 感应电流
    # Z_cav = R_Q / (1 / Q_t + 1j * ((omega_0 ** 2 - omega_cav ** 2) / (omega_0 * omega_cav)))  # 腔阻抗
    # abs_Z_cav = np.real(Z_cav)
    # theta_Z_cav = np.imag(Z_cav)
    # 感应电流的一次谐波I_ind_1
    # f1_t = np.array([I_ind(trange, z_i) * np.cos(omega_0 * trange) for trange in t])  # 离散取点
    # f2_t = np.array([I_ind(trange, z_i) * np.sin(omega_0 * trange) for trange in t])

    # F1 = 2 / T_0 * trapezoid(f1_t, [t], dx=0.01 * T_0)  # a1  # 离散求和代替连续积分数值求解方法
    # F2 = 2 / T_0 * trapezoid(f2_t, [t], dx=0.01 * T_0)  # b1

    # 感应电流一次谐波的幅值与相位
    # abs_I_ind_1 = np.sqrt(F1 ** 2 + F2 ** 2)
    # theta_I_ind_1 = np.arctan2(-F2, F1)  # arctan2求解范围的定义域[-pi,pi]
    # Vtg = abs_I_ind_1 * abs_Z_cav  # 腔压
    return Vtg * f(z) *np.exp(1j*omega_0 * t )  # 腔场|E_cav|


# 计算加速电场 Ez(t, z)
def get_E(t, z, z_i,  Vcav,zmin):
    E_spch_vals =   numpy.sum([sc_generator.get_Es(z, z_) for z_ in z_i],axis = 0)#E_spch(z, z_i, Q, a, b)
    E_cav_vals = E_cav(t, z,Vtg= Vcav)
    _f = scipy.interpolate.interp1d(z,  E_spch_vals + E_cav_vals)
    return    numpy.piecewise( z, [z>zmin,],[lambda z :_f (z),0] ) # 加速电场为空间电荷场与腔场之和


# 进行射频周期迭代计算，并根据收敛条件更新电场
def run_simulation(initial_pos,initial_vel,
                   Q, a, b, tolerance=1e-6, max_cycles=100):
    # 时间点
    t = np.linspace(0, T_0, t_steps)
    p0 = common.Ek_to_gamma_beta(50e3)  * C.m_e * C.c
    v0= common.p_to_v(p0)
    # dt = t[1]-  t[0]
    # 初始条件 (应当只设将进入谐振腔的电子圆盘)
    initial_conditions = np.array([[z0, p0] for z0 in np.linspace(z_min - v0 * T_0, z_max, z_steps)])

    # 存储每次计算的腔压，用于收敛判断
    Vtg_prev = [0j]
    converged = False
    cycle = 0

    while not converged and cycle < max_cycles:
        # 运动方程
        sol = solve_ivp(
            lorentz, t_span=(t[0], t[-1]), y0=initial_conditions.flatten(),
            method='RK45', t_eval=t, args=(initial_conditions[:, 0], Q, a, b,Vtg_prev[-1],z_min

                                        )  # 传入 z_i, Q, a, b
        )
        res =  sol.y.reshape((-1,*initial_conditions.shape))



        velocities =common.p_to_v( res[:,:,1])
        positions = res[:,:, 0]
        I_ind = np.sum(Q * velocities * f(positions),axis = 1)  # 感应电流
        Z_cav = R_Q / (1 / Q_t + 1j * ((omega_0 ** 2 - omega_cav ** 2) / (omega_0 * omega_cav)))  # 腔阻抗
        abs_Z_cav = np.real(Z_cav)
        theta_Z_cav = np.imag(Z_cav)
        # 感应电流的一次谐波I_ind_1
        f1_t = I_ind * np.cos(omega_0 * t)   # 离散取点
        f2_t = I_ind * np.sin(omega_0 * t)
        # plt.figure()
        # plt.plot(    t, I_ind)


        F1 = 2 / T_0 * trapezoid(f1_t, t, #dx=t[1]-t[0]
                                 )  # a1  # 离散求和代替连续积分数值求解方法
        F2 = 2 / T_0 * trapezoid(f2_t, t, #dx=t[1]-t[0]
                                 )  # b1
        Vcav = Z_cav * (F1 + 1j *F2 )
        Vtg_prev.append(Vcav)



        # # 计算空间电荷场 E_spch
        # z_i = sol.y[::2, :]  # 电子圆盘位置
        # E_spch_values = E_spch(np.linspace(z_min, z_max, 1000), z_i, Q, a, b)
        #
        # # 计算腔场 E_cav
        # E_cav_values = E_cav(t[-1], np.linspace(z_min, z_max, 1000), z_i)  # 在最后一个时刻计算腔场
        #
        # # 线性插值，确保电场长度一致
        # interp_E_spch = interp1d(np.linspace(z_min, z_max, 1000), E_spch_values, kind='linear', fill_value='extrapolate')
        # interp_E_cav = interp1d(np.linspace(z_min, z_max, 1000), E_cav_values, kind='linear', fill_value='extrapolate')
        #
        # E_spch_values_interp = interp_E_spch(np.linspace(z_min, z_max, 500))
        # E_cav_values_interp = interp_E_cav(np.linspace(z_min, z_max, 500))
        #
        # # 计算加速电场 Ez(t, z)
        # Ez_values = E_spch_values_interp + E_cav_values_interp
        #
        # # 计算腔压 Vtg
        # Vtg = np.sum(E_cav_values_interp)

        # 判断收敛：检查腔压的变化是否小于容差
        if Vtg_prev is not None:
            Vtg_diff = np.abs(Vcav - Vtg_prev[-2])
            if Vtg_diff < tolerance  *  np.abs(Vtg_prev[-1]):
                converged = True

        # Vtg_prev = Vtg
        cycle += 1
        logger.info("cycle = %d, Vcav = %.2e"%(cycle,numpy.abs(Vcav)))


    # 输出最终结果
    print(f"Simulation converged after {cycle} RF cycles.")
    return sol, Vtg_prev

if __name__ == '__main__':
    plt.ion()


    # 主程序执行
    ts = numpy.linspace(0, 1/f0, 200)
    p0 = common.Ek_to_gamma_beta(50e3) * C.m_e * C.c
    v0 = common.p_to_v(p0)
    zs = np.linspace(0 - v0 * T_0, 20e-3, 100)

    plt.figure()
    plt.plot(zs, f(zs))
    Nrings = len(zs)


    # dt = t[1]-  t[0]
    # 初始条件 (应当只设将进入谐振腔的电子圆盘)



    sol, Vtg_prev = run_simulation(z_min, z_max, t_steps, z_steps, Q, a, b)
    plt.figure()
    res = sol.y.reshape((-1,z_steps, 2))
    plt.plot(sol.t, res[:, :,0] )

    plt.figure()
    plt.plot(Vtg_prev)
    aaaaaaa

    # 绘制结果
    plt.plot(np.linspace(z_min, z_max, 500), numpy.abs(Vtg_prev))
    plt.xlabel("Position (m)")
    plt.ylabel("Electric Field (V/m)")
    plt.title("Electric Field Distribution after Convergence")
    plt.show()